import jieba
import numpy as np
import matplotlib.pyplot as plt

corpus = ["我特别特别喜欢看电影",
          "这部电影真的是很好看的电影",
          "今天天气真好使难得的好天气",
          "我今天去看了一部电影",
          "电影院的电影都很好看"]

# 对句子分词
corpus_tokenized = [list(jieba.cut(sentence)) for sentence in corpus]
print(corpus_tokenized)

# 创建词汇表
word_dict = {}
for sentence in corpus_tokenized:
    for word in sentence:
        if word not in word_dict:
            # 分配当前词汇表索引
            word_dict[word] = len(word_dict)
print(f"\n{word_dict}")

# 根据词汇表将句子转换为词袋表示
bow_vectors = []
for sentence in corpus_tokenized:
    # 初始化全0向量
    sentence_vector = [0] * len(word_dict)
    for word in sentence:
        # 给对应词索引位置上的数字加1, 表示该词在当前句子中出现了1次
        sentence_vector[word_dict[word]] += 1
    # 将当前句子的词袋向量添加到向量列表中
    bow_vectors.append(sentence_vector)
print("\n词袋表示: ", bow_vectors)


# 定义余弦相似度函数
def consine_similarity(vec1, vec2):
    # 计算向量1 和 向量2 的点积
    dot_product = np.dot(vec1, vec2)
    # 向量1的范数
    norm_a = np.linalg.norm(vec1)
    # 向量1的范数
    norm_b = np.linalg.norm(vec2)
    # 余弦相似度
    return dot_product / (norm_a * norm_b)

# 初始化全0矩阵，用于存储余弦相似度
similarity_matrix = np.zeros((len(corpus), len(corpus)))

# 计算每两个句子之间的余弦相似度

for i in range(len(corpus)):
    for j in range(len(corpus)):
        similarity_matrix[i][j] = consine_similarity(bow_vectors[i], bow_vectors[j])


# 用于可视化余弦相似度
# 设定字体样式
plt.rcParams["font.family"] = ["SimHei"]
# 设定无衬线字体样式
plt.rcParams["font.sans-serif"] = ["SimHei"]
# 用来正常显示负号
plt.rcParams["axes.unicode_minus"] = False
# 创建一个绘图对象
fig, ax = plt.subplots()
# 使用matshow函数绘制余弦相似度矩阵，颜色使用蓝色
cax = ax.matshow(similarity_matrix, cmap=plt.cm.Blues)
# 条形图颜色映射
fig.colorbar(cax)
# X轴刻度
ax.set_xticks(range(len(corpus)))
# Y轴刻度
ax.set_yticks(range(len(corpus)))
# 刻度标签
ax.set_xticklabels(corpus, rotation=45, ha='left')
# 刻度标签为原始句子
ax.set_yticklabels(corpus)
# 显示图形
plt.show()